import gradio as gr
from PIL import Image
import io
import base64
from pathlib import Path
import logging
import re
from typing import Dict, List, Tuple
import os
import time
import warnings
from deepseek_api import get_deepseek_response

class SimpleChatUI:
    """简化版聊天界面"""

    def __init__(self):
        self.theme = gr.themes.Soft(
            primary_hue="blue",
            secondary_hue="blue"
        )
        self.history = []
        self.static_dir = Path("static")
        self.static_dir.mkdir(parents=True, exist_ok=True)
        self.paused = False
        self.interrupt_state = False  # 新增：按钮状态（发送状态或打断状态）
        current_dir = os.path.dirname(__file__)
        self.user_config = os.path.join(current_dir, r'user.png')
        self.agent_config_list = os.path.join(current_dir, r'little_fat.png')

    def create_interface(self):
        with gr.Blocks(
                theme=self.theme,
                css="""
            /* === 核心布局 === */
            .root-container {
                /* 设置绝对最小宽度 */
                min-width: 800px !important;
                width: 100%;
                /* 强制不换行 */
                display: flex;
                flex-wrap: nowrap !important;
            }

            /* === 两栏比例 === */
            .chat-panel {
                flex: 4;
                /* 最小宽度 = 800px * 0.75 = 600px */
                min-width: 600px;  
            }

            .info-panel {
                flex: 1;
                /* 最小宽度 = 800px * 0.25 = 200px */
                min-width: 200px;  
            }

            /* === 组件保护 === */
            /* 聊天框 */
            .chatbot {
                height: 70vh;
                min-height: 300px;
            }

            /* 输入框 */
            .textbox {
                min-width: 200px !important;
                outline: none !important; /* 移除聚焦时的蓝色边框 */
                box-shadow: none !important; /* 移除聚焦时的阴影效果 */
            }
            

            /* 按钮 */
            .btn {
                min-width: 80px !important;
                width: auto !important;
            }

            /* 中断按钮样式 */
            .interrupt-btn {
                min-width: 80px !important;
                width: auto !important;
            }

            /* 禁用按钮样式 */
            .interrupt-btn:disabled {
                background-color: #cccccc;
                cursor: not-allowed;
            }

            /* === 响应式保护 === */
            @media (max-width: 800px) {
                /* 小于800px时启用强制缩放 */
                html {
                    zoom: calc(100vw / 800);
                }
            }
            .generating {
                border: none !important;
                box-shadow: none !important;
            }
            .wrap.generating {
                display: none !important;
            }
            """
        ) as interface:
            # 主容器（强制最小宽度800px）
            with gr.Row(elem_classes="root-container"):
                # 左侧聊天区 (75%)
                with gr.Column(elem_classes="chat-panel", scale=4):
                    self.chatbot = gr.Chatbot(elem_classes="chatbot", height=900,
                                              avatar_images=[self.user_config, self.agent_config_list])

                    with gr.Row():
                        self.msg = gr.Textbox(
                            placeholder="输入消息...",
                            elem_classes="textbox",
                            scale=4
                        )

                    with gr.Row():
                        self.clear_btn = gr.Button("清空", elem_classes="btn")
                        self.submit_btn = gr.Button("发送", variant="primary", elem_classes="btn")

                        # 新增中断按钮，使用自定义 CSS 类
                        self.interrupt_btn = gr.Button("中断", variant="secondary", elem_classes="interrupt-btn")

                # 右侧信息区 (25%)
                with gr.Column(elem_classes="info-panel", scale=1):
                    # 读取本地图片并转换为 Base64 编码
                    img_path = Path(r"icon.jpg")
                    if img_path.exists():
                        with open(img_path, "rb") as img_file:
                            img_data = img_file.read()
                            img_base64 = base64.b64encode(img_data).decode("utf-8")
                    else:
                        img_base64 = ""  # 如果图片不存在，设置为空

                    gr.Markdown(f"![图标](data:image/jpeg;base64,{img_base64})")
                    gr.Markdown("## 代理信息")
                    gr.Markdown("**名称**: 简化版助手")

                self._setup_event_handlers()

        return interface

    def _setup_event_handlers(self):
        """设置事件处理器"""
        self.msg.submit(
            self._handle_submit,
            [self.msg],
            [self.msg, self.chatbot, self.submit_btn, self.interrupt_btn],
            api_name="submit"  # 添加api_name避免警告
        )

        self.submit_btn.click(
            self._handle_submit,
            [self.msg],
            [self.msg, self.chatbot, self.submit_btn, self.interrupt_btn],
            api_name="submit"
        )

        self.clear_btn.click(
            self._clear_chat,
            outputs=[self.chatbot]
        )

        # 新增中断按钮事件
        self.interrupt_btn.click(
            fn=self._toggle_interrupt,
            outputs=[self.interrupt_btn, self.submit_btn],
            api_name="toggle_interrupt"
        )

    def _toggle_interrupt(self):
        """切换中断状态"""
        if self.interrupt_state:
            self.interrupt_state = False
            return gr.update(interactive=False), gr.update(interactive=True)
        else:
            self.interrupt_state = True
            return gr.update(interactive=False), gr.update(interactive=True)

    def _handle_submit(self, message: str):
        """处理提交事件（支持真正的打断功能）"""
        if not message.strip() and not self.interrupt_state:
            # 如果消息为空且未处于中断状态，直接返回当前状态
            yield "", self.chatbot, gr.update(), gr.update()
            return

        if self.interrupt_state:
            yield "", self._format_chat_history(), gr.update(interactive=True), gr.update(
                interactive=False)
            return

        # 添加用户消息
        self._add_to_history("user", {"text": message})

        # 初始化空回复
        self._add_to_history("assistant", {"text": ""})

        # 创建生成器实例
        # response_generator = self._generate_response(message)
        response_generator = get_deepseek_response(message)
        

        try:
            while True:
                if self.interrupt_state:  # 检查中断状态
                    raise GeneratorExit("用户主动中断")

                # 获取下一块内容
                chunk = next(response_generator)

                # 检查历史是否为空（防止清空后继续访问）
                if not self.history:
                    break

                # 更新最后一条助手消息
                self.history[-1]["content"]["text"] += chunk

                # 返回更新后的聊天记录
                yield "", self._format_chat_history(), gr.update(interactive=False), gr.update(interactive=True)
                time.sleep(0.05)

        except (GeneratorExit, StopIteration):
            # 正常结束或被中断
            pass

        finally:
            # 最终状态
            yield "", self._format_chat_history(), gr.update(interactive=True), gr.update(interactive=False)
            self.interrupt_state = False  # 重置中断状态

    # def _generate_response(self, message: str):
    #     """流式生成助手回复（支持中断）"""
    #     if not self.history:
    #         return

    #     text_content = self._read_text_file("text.txt")
    #     code_content = self._read_text_file("code.txt")

    #     # 先流式输出文本和代码部分
    #     text_response = f"{text_content}\n\n```python\n{code_content}\n```"
    #     for char in text_response:
    #         if self.interrupt_state:  # 检查中断状态
    #             return
    #         yield char
    #         time.sleep(0.05)  # 添加延迟模拟流式效果

    #     # 文本流式输出完成后，一次性追加图片
    #     img_html = self._get_img_html("icon.jpg")
    #     if img_html:
    #         yield img_html






    def _get_img_html(self, file_path: str) -> str:
        """获取图片的HTML格式"""
        try:
            img_path = Path(file_path)
            if not img_path.exists():
                return "<p>[图片文件不存在]</p>"

            with Image.open(img_path) as img:
                img.thumbnail((400, 400))
                buffered = io.BytesIO()
                img.save(buffered, format="PNG")
                img_str = base64.b64encode(buffered.getvalue()).decode()
                return f'<img src="data:image/png;base64,{img_str}" style="max-width:400px; margin-top:20px;">'
        except Exception as e:
            logging.error(f"读取图片失败: {str(e)}")
            return f"<p>[图片加载失败: {str(e)}]</p>"

    def _read_text_file(self, file_path: str) -> str:
        """读取文本文件内容"""
        try:
            with open(file_path, "r", encoding="utf-8") as f:
                return f.read()
        except Exception as e:
            logging.error(f"读取文件失败: {str(e)}")
            return ""

    def _format_chat_history(self) -> List[Tuple[str, str]]:
        """格式化聊天历史"""
        formatted_history = []
        for i in range(0, len(self.history), 2):
            if i + 1 < len(self.history):
                user_msg = self.history[i]["content"]["text"]
                assistant_msg = self.history[i + 1]
                formatted_content = self._format_display_content(assistant_msg["content"])
                formatted_history.append((user_msg, formatted_content))
        return formatted_history

    def _format_display_content(self, content: Dict) -> str:
        """格式化显示内容（支持HTML和Markdown混合）"""
        if "text" not in content:
            return ""

        text = content["text"]

        # 处理代码块
        code_blocks = re.findall(r'```python\n(.*?)\n```', text, re.DOTALL)
        for code in code_blocks:
            folded_code = f"<details class='code-block'><summary>点击展开代码</summary><pre><code>{code}</code></pre></details>"
            text = text.replace(f"```python\n{code}\n```", folded_code)

        # 保留HTML标签
        return text

    def _add_to_history(self, role: str, content: Dict):
        """添加对话记录"""
        self.history.append({
            "role": role,
            "content": content
        })

    def _clear_chat(self):
        """清空对话"""
        self.history.clear()
        self.interrupt_state = False  # 重置中断状态
        return self._format_chat_history()  # 返回格式化的聊天历史（此时为空）


if __name__ == "__main__":
    app = SimpleChatUI()
    interface = app.create_interface()
    interface.queue()  # 启用队列支持流式输出
    interface.launch(
        server_name="127.0.0.1",
        server_port=7861,
        share=False
    )